import random
from configs import split_config


def yield_read(file):
    with open(file, "r") as h_file:
        while True:
            line = h_file.readline()
            if line:
                yield line
            else:
                break


def read(file):
    with open(file, "r") as h_file:
        return h_file.readlines()


if not split_config.shuffle:
    src_data = yield_read(split_config.src_file)
    tgt_data = yield_read(split_config.tgt_file)
else:
    src_data = read(split_config.src_file)
    random.shuffle(src_data)
    tgt_data = read(split_config.tgt_file)
    random.shuffle(tgt_data)

index = 0

src_train_file = open(split_config.src_train_file, "w+")
src_test_file = open(split_config.src_test_file, "w+")
src_valid_file = open(split_config.src_valid_file, "w+")

tgt_train_file = open(split_config.tgt_train_file, "w+")
tgt_test_file = open(split_config.tgt_test_file, "w+")
tgt_valid_file = open(split_config.tgt_valid_file, "w+")

for src_item, tgt_item in zip(src_data, tgt_data):
    r = random.random()
    if r < split_config.test_radio:
        src_test_file.write(src_item)
        tgt_test_file.write(tgt_item)
    else:
        src_train_file.write(src_item)
        tgt_train_file.write(tgt_item)

    if index < split_config.valid_num:
        src_valid_file.write(src_item)
        tgt_valid_file.write(tgt_item)

    index = index + 1
